#include <Dynamics/Fem/fem_cpu.h>
#include <iostream>
#include <Utilities/glm_ext.h>
#include <Dynamics/Fem/shape_function.h>
using namespace std;

namespace PhysLeo {

template<typename T>
FemCpu<T>::FemCpu(): ForceModule<T>(),
                                        quadrature_points_(1),
                                        fem_material_type_(FemMaterialType::NEOHOOKEAN),
                                        homogeneous_(true),
                                        ele_type_(EleType::NONE)
{
}

template<typename T>
void FemCpu<T>::allocateResource(SimData<T> & simData)
{
    //buffer data
    ptr_rest_pos_ = simData.ptr_rest_pos_;
    ptr_pos_ = simData.ptr_pos_;
    ptr_force_ = simData.ptr_force_;
    ptr_ele_ = simData.ptr_ele_;
    ptr_ele_volume_ = simData.ptr_ele_volume_;
    ptr_ele_material_ = simData.ptr_ele_material_;

    //property data
    ele_type_ = simData.ele_type_;
    fem_material_type_ = simData.fem_material_type_;
    quadrature_points_ = simData.quadrature_points_;
    homogeneous_ = simData.homogeneous_;

    //compability check for the simulation data
    if (quadrature_points_ > 1 && (ele_type_ == EleType::TET || ele_type_ == EleType::TRI)) 
    {
        cout << "FEM Model: quadrature points for element type TET4 and TRI3 must be 1, because the strain field in these two element is constant!" << endl;
        cout << "system has automatically changed the quadraturePoints to 1." << endl;
        quadrature_points_ = 1;
    }

    if(quadrature_points_ >=3 )
    {
        cout << "quadrature points greater than 2 is not supported now! reduce it to 2" << endl;
        quadrature_points_ = 2;
    }
    
    const int ele_vert_num = eleVertNumFromEleType(ele_type_);
    int ele_num = ptr_ele_->size() / ele_vert_num;
    if (ele_num != ptr_ele_volume_->size())
    {
        cout << "element data doesn't compatible with element type!" << endl;
        exit(1);
    }
    
    //allocate cuda memory for shape matrix according to eleType and quadraturePoints
    //compute shapeMatrix
    switch (ele_type_)
    {
    case PhysLeo::EleType::NONE:
        cout << "Error:element type is None!" << endl;
        exit(1);
        break;
    case PhysLeo::EleType::TRI:
        cout << "Error: FEM haven't been implemented for element type TRI" << endl;
        exit(1);
        break;
    case PhysLeo::EleType::QUAD:
        cout << "Error: FEM haven't been implemented for element type QUAD" << endl;
        exit(1);
        break;
    case PhysLeo::EleType::TET:
        tetShapeMatrix();
        break;
    case PhysLeo::EleType::HEX:
        cubicShapeMatrix();
        break;
    default:
        cout << "Error: element type is unidentified!" << endl;
        exit(1);
        break;
    }
}

template<typename T>
T FemCpu<T>::deformEnergy(std::shared_ptr<BufferData<glm::tvec3<T>>> ptr_pos)
{
    return T();
}

template<typename T>
void FemCpu<T>::internalForce(std::shared_ptr<BufferData<glm::tvec3<T>>> ptr_pos, std::shared_ptr<BufferData<glm::tvec3<T>>> ptr_force)
{
    if (ele_type_ == EleType::HEX) {
        cubicInternalForce(ptr_pos, ptr_force);
    }
    else if (ele_type_ == EleType::TET) {
        tetInternalForce(ptr_pos, ptr_force);
    }
}

template<typename T>
void FemCpu<T>::cubicShapeMatrix()
{
    //gaussian quadrature points integration information
    int sample_num[2] = { 1,8 };
    auto a = static_cast<T>(0.57735);
    glm::tvec3<T> one_points_pos[1] = { glm::tvec3<T>(0,0,0) };
    glm::tvec3<T> two_points_pos[8] = { glm::tvec3<T>(-a,-a,-a), glm::tvec3<T>(-a,-a,a), glm::tvec3<T>(-a,a,-a),  glm::tvec3<T>(-a,a,a),
        glm::tvec3<T>(a,-a,-a), glm::tvec3<T>(a,-a,a), glm::tvec3<T>(a,a,-a),glm::tvec3<T>(a,a,a) };
    glm::tvec3<T>* sample_point[2] = { one_points_pos, two_points_pos };

    const int ele_num = ptr_ele_volume_.get()->size();
    ptr_shape_matrix_ = make_shared<BufferData<T>>(3 * 8 * ele_num*sample_num[quadrature_points_ - 1]);
    ptr_shape_matrix_->allocateCpu();

    auto mat_ptr = reinterpret_cast<glm::tmat8x3<T>*>(ptr_shape_matrix_->ptrCpu());
    auto rest_pos_ptr = ptr_rest_pos_->ptrCpu();

    for (auto ele_id = 0; ele_id < ele_num; ++ele_id)
    {
        glm::tmat8x3<T> rest_shape;
        for (auto i = 0; i < 8; ++i) {
            auto vert_index = (ptr_ele_->ptrCpu())[ele_id * 8 + i];
            rest_shape[i] = rest_pos_ptr[vert_index];
        }

        for (auto i = 0; i < sample_num[quadrature_points_-1]; ++i) 
        {
            glm::tmat8x3<T> s;
            cubicShapeDerivatives(sample_point[quadrature_points_-1][i], s);
            auto matrix_dm = rest_shape*s;
            mat_ptr[ele_id*sample_num[quadrature_points_ - 1] + i] = s * glm::inverse(matrix_dm);
        }
    }
}

template<typename T>
void FemCpu<T>::tetShapeMatrix()
{
    const int ele_num = ptr_ele_volume_->size();
    ptr_shape_matrix_ = make_shared<BufferData<T>>(3 * 4 * ele_num);
    ptr_shape_matrix_->allocateCpu();

    auto mat_ptr = reinterpret_cast<glm::tmat3x4<T>*>(ptr_shape_matrix_->ptrCpu());
    auto rest_pos_ptr = ptr_rest_pos_->ptrCpu();

    for (auto ele_id = 0; ele_id < ele_num; ++ele_id)
    {
        glm::tmat4x3<T> rest_shape;  //3行4列矩阵
        for (auto i = 0; i < 4; ++i) {
            auto vert_index = (ptr_ele_->ptrCpu())[4 * ele_id + i];
            rest_shape[i] = rest_pos_ptr[vert_index];
        }

        //实际上是 4行3列的矩阵
        glm::tmat3x4<T> s(-1.0, 1.0, 0.0, 0.0, -1.0, 0.0, 1.0, 0.0, -1.0, 0.0, 0.0, 1.0);
        mat_ptr[ele_id] = s * glm::inverse(rest_shape * s);
    }
}

template<typename T>
void FemCpu<T>::cubicInternalForce(std::shared_ptr<BufferData<glm::tvec3<T>>> ptr_pos, std::shared_ptr<BufferData<glm::tvec3<T>>> ptr_force)
{
    int sample_num[2] = { 1,8 };
    T one_point_weight[1] = { 1.0 };
    T two_point_weight[8] = { 0.125,0.125,0.125,0.125, 0.125,0.125, 0.125,0.125 };
    T* sample_weight[2] = { one_point_weight,two_point_weight };

    //material
    const int num_material_parameter = FemMaterial<T>::numMaterialParams(fem_material_type_);

    auto mat_ptr = reinterpret_cast<glm::tmat8x3<T>*>(ptr_shape_matrix_->ptrCpu());
    auto pos_ptr = ptr_pos->ptrCpu();
    auto force_ptr = ptr_force->ptrCpu();
    const auto material_ptr = ptr_ele_material_->ptrCpu();

    const int ele_num = ptr_ele_volume_->size();
    for (auto ele_id = 0; ele_id < ele_num; ++ele_id)
    {
        glm::tmat8x3<T> current_shape;
        for (auto i = 0; i < 8; ++i)
        {
           auto vert_index = (ptr_ele_->ptrCpu())[ele_id * 8 + i];
            current_shape[i] = pos_ptr[vert_index];
        }
        auto temp_ptr = material_ptr + (homogeneous_ ? 0 : ele_id * num_material_parameter);

        //force
        glm::tmat8x3<T> total_force(0.0);
        for (auto i = 0; i<sample_num[quadrature_points_-1]; ++i)
        {
            glm::tmat8x3<T> rest_shape = mat_ptr[ele_id*sample_num[quadrature_points_ - 1] + i];
            glm::tmat3x3<T> deform_gradient = current_shape * rest_shape;

            //solve invertible problem
            if (glm::determinant(deform_gradient) < 0)
            {
                glm::tvec3<T> col = deform_gradient[0];
                deform_gradient[0] = deform_gradient[1];
                deform_gradient[1] = col;
            }

            glm::tmat3x3<T> stress = FemMaterial<T>::firstPiolaStress(deform_gradient, temp_ptr, fem_material_type_);
            total_force = total_force + static_cast<T>(-1) * sample_weight[quadrature_points_ - 1][i] * (ptr_ele_volume_->ptrCpu())[ele_id] * stress * rest_shape;
        }

        //add
        for (auto i = 0; i < 8; ++i) {
            auto vert_index = (ptr_ele_->ptrCpu())[ele_id * 8 + i];
            force_ptr[vert_index] += total_force[i];
        }
    }
}

template<typename T>
void FemCpu<T>::tetInternalForce(std::shared_ptr<BufferData<glm::tvec3<T>>> ptr_pos, std::shared_ptr<BufferData<glm::tvec3<T>>> ptr_force)
{
    //material
    const int num_material_parameter = FemMaterial<T>::numMaterialParams(fem_material_type_);

    auto mat_ptr = reinterpret_cast<glm::tmat3x4<T>*>(ptr_shape_matrix_->ptrCpu());
    auto pos_ptr = ptr_pos->ptrCpu();
    auto force_ptr = ptr_force->ptrCpu();
    const auto material_ptr = ptr_ele_material_->ptrCpu();

    const int ele_num = ptr_ele_volume_->size();
    for (auto ele_id = 0; ele_id < ele_num; ++ele_id)
    {
        glm::tmat4x3<T> current_shape;  //3行4列矩阵
        for (auto i = 0; i < 4; ++i) {
            auto vert_index = (ptr_ele_->ptrCpu())[ele_id * 4 + i];
            current_shape[i] = pos_ptr[vert_index];
        }

        auto temp_ptr = material_ptr + (homogeneous_ ? 0 : ele_id * num_material_parameter);

        glm::tmat3x4<T> rest_shape = mat_ptr[ele_id];
        glm::tmat3x3<T> deform_gradient = current_shape * rest_shape;

        //solve invertible problem
        if (glm::determinant(deform_gradient) < 0)
        {
            glm::tvec3<T> col = deform_gradient[0];
            deform_gradient[0] = deform_gradient[1];
            deform_gradient[1] = col;
        }

        glm::tmat3x3<T> stress = FemMaterial<T>::firstPiolaStress(deform_gradient, temp_ptr, fem_material_type_);
        glm::tmat4x3<T> total_force = -(ptr_ele_volume_->ptrCpu())[ele_id] * stress * glm::transpose(rest_shape);

        for (auto i = 0; i < 4; ++i) {
            auto vert_index = (ptr_ele_->ptrCpu())[ele_id * 4 + i];
            force_ptr[vert_index] += total_force[i];
        }
    }
}

template class FemCpu<float>;
template class FemCpu<double>;

}